package zkp

import (
	"fmt"
	"log"
	"os"

	"github.com/consensys/gnark/backend/groth16"
	"github.com/consensys/gnark/constraint"
	"github.com/consensys/gnark/frontend"
	"github.com/consensys/gnark/frontend/cs/r1cs"
)

const (
	// 客户端只需要 proving key，由开发团队提供
	ProvingKeyFile = "auth_proving.key"
)

// KeyManager 密钥管理器（客户端版本）
// 只加载 proving key，不生成密钥
type KeyManager struct {
	provingKey groth16.ProvingKey
	ccs        constraint.ConstraintSystem
}

// NewKeyManager 创建密钥管理器（客户端版本）
// 只加载 proving key，不生成密钥
// Proving key 应该由开发团队通过 Trusted Setup 生成并分发
func NewKeyManager() (*KeyManager, error) {
	km := &KeyManager{}

	// 检查 proving key 是否存在
	if !fileExists(ProvingKeyFile) {
		return nil, fmt.Errorf(`
❌ 错误: %s 不存在

客户端不应该生成密钥！Proving key 应该由开发团队统一生成并分发。

解决方案:
1. 从开发团队获取 %s
2. 将文件放置在项目根目录（与 main.go 同级）
3. 重新运行程序

参考文档: ZKP完整流程说明.md
`, ProvingKeyFile, ProvingKeyFile)
	}

	log.Println("加载 proving key...")
	if err := km.loadKeys(); err != nil {
		return nil, fmt.Errorf("加载密钥失败: %v", err)
	}
	log.Println("✅ Proving key 加载成功")

	return km, nil
}

// loadKeys 从文件加载密钥（客户端版本）
// 只加载 proving key 和编译电路（生成 ccs）
func (km *KeyManager) loadKeys() error {
	// 加载 proving key
	pk, err := loadProvingKey(ProvingKeyFile)
	if err != nil {
		return fmt.Errorf("加载 proving key 失败: %v", err)
	}

	// 编译电路以获取 constraint system (ccs)
	// 这是生成证明时必需的，但不涉及密钥生成
	circuit := &AuthCircuit{}
	ccs, err := frontend.Compile(GetCurve().ScalarField(), r1cs.NewBuilder, circuit)
	if err != nil {
		return fmt.Errorf("编译电路失败: %v", err)
	}

	km.provingKey = pk
	km.ccs = ccs

	return nil
}

// GetProver 获取证明生成器
func (km *KeyManager) GetProver() *Prover {
	return NewProver(km.provingKey, km.ccs)
}

// === 辅助函数 ===

// fileExists 检查文件是否存在
func fileExists(filename string) bool {
	_, err := os.Stat(filename)
	return err == nil
}

// loadProvingKey 从文件加载 proving key
func loadProvingKey(filename string) (groth16.ProvingKey, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	pk := groth16.NewProvingKey(GetCurve())
	_, err = pk.ReadFrom(file)
	return pk, err
}
